%%%-------------------------------------------------------------------
%%% @author mirahs(2892727493@qq.com)
%%% @doc
%%%     mnesia 接口封装
%%% @end
%%%-------------------------------------------------------------------
-module(db_mnesia).

-export([
    write/2    %% 写入
    ,read/2     %% 读取数据，根据主键
    ,read/3     %% 读取数据，根据索引
    ,select/2   %% 读取数据，使用match pattern
    ,all_keys/1 %% 读取表的所有key
    ,delete/2   %% 删除，根据主键
    ,clear/1    %% 清除表中所有数据
    ,foldl/2    %% mnesia foldl
    ,backup/1   %% 备份
    ,restore/1  %% 恢复
    ,restore/2  %% 恢复
]).

-include("common.hrl").
-include("mnesia.hrl").
-include("env.hrl").


%% @spec write(Tab, Record) -> {ok, Num} | {error, Reason}
%% Tab = atom() 表名
%% Record = term() 数据记录
%% Num = integer() 成功的条数
%% Reason = term() 错误信息
%% 写入记录
write(Tab, Record) ->
    Fun = case mnesia:table_info(Tab, frag_properties) of
              [] ->
                  fun(T, R) -> mnesia:dirty_write(T, R) end;
              FragCfg when is_list(FragCfg) ->
                  F = fun(T, R) -> mnesia:write(T, R, write) end,
                  fun(T, R) -> mnesia:activity(async_dirty, F, [T, R], mnesia_frag) end
          end,
    case catch Fun(Tab, Record) of
        ok -> {ok, 1};
        {'EXIT', {aborted, Reason}} ->
            {error, Reason}
    end.


%% @spec read(Tab, Key) -> {ok, Vl} | {error, Reason}
%% Tab = atom() 表名
%% Key = term() 主键值
%% Vl = [#term{} | ..] 结果列表
%% Reason = term() 错误信息
%% 根据键值，读取数据
read(Tab, Key) ->
    Fun = case mnesia:table_info(Tab, frag_properties) of
              [] ->
                  fun(T, K) -> mnesia:dirty_read(T, K) end;
              FragCfg when is_list(FragCfg) ->
                  F = fun({T, K}) -> mnesia:read(T, K) end,
                  fun(T, K) -> mnesia:activity(sync_dirty, F, [{T, K}], mnesia_frag) end
          end,
    case catch Fun(Tab, Key) of
        Vl when is_list(Vl) ->
            {ok, Vl};
        {'EXIT', {aborted, Reason}} ->
            {error, Reason}
    end.

%% @spec read(Tab, Pos, Key) -> {ok, Vl} | {error, Reason}
%% Tab = atom() 表名
%% Pos = integer() 位置
%% Key = term() 索引值
%% Vl = [#term{} | ..] 结果列表
%% Reason = term() 错误信息
%% 根据索引值，读取数据
%% 注意: Pos必须是索引，否则报错
read(Tab, Pos, Key) ->
    Fun = case mnesia:table_info(Tab, frag_properties) of
              [] ->
                  fun(T, K, P) -> mnesia:dirty_index_read(T, K, P) end;
              FragCfg when is_list(FragCfg) ->
                  F = fun({T, K, P}) -> mnesia:index_read(T, K, P) end,
                  fun(T, K, P) -> mnesia:activity(sync_dirty, F, [{T, K, P}], mnesia_frag) end
          end,
    case catch Fun(Tab, Key, Pos) of
        Vl when is_list(Vl) ->
            {ok, Vl};
        {'EXIT', {aborted, Reason}} ->
            {error, Reason}
    end.

%% @spec select(Tab, Pattern) -> {ok, Vl} | {error, Reason}
%% Tab = atom() 表名
%% Pattern = match_pattern() | match_spec()
%% Vl = [#term{} | ..] 结果列表
%% Reason = term() 错误信息
%% 根据match_pattern，读取数据
select(Tab, Pattern) when is_tuple(Pattern) ->
    Fun = case mnesia:table_info(Tab, frag_properties) of
              [] ->
                  fun(T, Pat) -> mnesia:dirty_match_object(T, Pat) end;
              FragCfg when is_list(FragCfg) ->
                  F = fun({T, Pat}) -> mnesia:match_object(T, Pat, none) end,
                  fun(T, Pat) -> mnesia:activity(sync_dirty, F, [{T, Pat}], mnesia_frag) end
          end,
    case catch Fun(Tab, Pattern) of
        Vl when is_list(Vl) ->
            {ok, Vl};
        {'EXIT', {aborted, Reason}} ->
            {error, Reason}
    end;
select(Tab, MatchSpec) when is_list(MatchSpec) ->
    Fun = case mnesia:table_info(Tab, frag_properties) of
              [] ->
                  fun(T, MS) -> mnesia:dirty_select(T, MS) end;
              FragCfg when is_list(FragCfg) ->
                  F = fun({T, MS}) -> mnesia:select(T, MS, read) end,
                  fun(T, MS) -> mnesia:activity(sync_dirty, F, [{T, MS}], mnesia_frag) end
          end,
    case catch Fun(Tab, MatchSpec) of
        Vl when is_list(Vl) ->
            {ok, Vl};
        {'EXIT', {aborted, Reason}} ->
            {error, Reason}
    end.

%% @spec all_keys(Tab) -> {ok, Vl} | {error, Reason}
%% Tab = atom() 表名
%% Vl = [#term{} | ..] 结果列表
%% Reason = term() 错误信息
%% 读取表的所有keys
all_keys(Tab) ->
    Fun = case mnesia:table_info(Tab, frag_properties) of
              [] ->
                  fun(T) -> mnesia:dirty_all_keys(T) end;
              FragCfg when is_list(FragCfg) ->
                  F = fun(T) -> mnesia:all_keys(T) end,
                  fun(T) -> mnesia:activity(sync_dirty, F, [T], mnesia_frag) end
          end,
    case catch Fun(Tab) of
        Vl when is_list(Vl) ->
            {ok, Vl};
        {'EXIT', {aborted, Reason}} ->
            {error, Reason}
    end.

%% @spec delete(Tab, Key) -> {ok, Num} | {error, Reason}
%% Tab = atom() 表名
%% Key = term() 主键值
%% Num = integer() 删除的数量
%% Reason = term() 错误信息
%% 根据Key, 删除数据
delete(Tab, Key) ->
    Fun = case mnesia:table_info(Tab, frag_properties) of
              [] ->
                  fun(T, K) -> mnesia:dirty_delete(T, K) end;
              FragCfg when is_list(FragCfg) ->
                  F = fun({T, K}) -> mnesia:delete(T, K, write) end,
                  fun(T, K) -> mnesia:activity(async_dirty, F, [{T, K}], mnesia_frag) end
          end,
    case catch Fun(Tab, Key) of
        ok -> {ok, 1};
        {'EXIT', {aborted, Reason}} ->
            {error, Reason}
    end.

%% @spec clear(Tab) -> ok | {error, Reason}
%% Tab = atom() 表名
%% Reason = term() 错误信息
%% 清除表数据
clear(Tab) ->
    Fun = case mnesia:table_info(Tab, frag_properties) of
              [] ->
                  fun(T) -> mnesia:clear_table(T) end;
              FragCfg when is_list(FragCfg) ->
                  F = fun(T) -> mnesia:clear_table(T) end,
                  fun(T) -> mnesia:activity(sync_dirty, F, [T], mnesia_frag) end
          end,
    case Fun(Tab) of
        {atomic, ok} -> ok;
        {aborted, Reason} ->
            {error, Reason}
    end.

%% @spec foldl(Fun, Tab) -> ok | {error, Reason}
%% Fun = fun 函数
%% Tab = atom() 表名
%% Reason = term() 错误信息
%% foldl
foldl(Fun, Tab) ->
    Fun1 = case mnesia:table_info(Tab, frag_properties) of
               [] ->
                   fun(F, T) ->  mnesia:activity(sync_dirty, fun mnesia:foldl/3, [F, [], T], mnesia) end;
               FragCfg when is_list(FragCfg) ->
                   fun(F, T) -> mnesia:activity(sync_dirty, fun mnesia:foldl/3, [F, [], T], mnesia_frag) end
           end,
    case Fun1(Fun, Tab) of
        [] -> ok;
        Other ->
            ?ERR("mnesia error:~w", [Other]),
            {error, Other}
    end.

%% @spec backup(FileName) -> ok | {error, Reason}
%% Filename = string()
%% Reason = term()
%% 备份整个数据库
backup(Path) ->
    mnesia:backup(Path).

%% @spec backup(FileBasename) -> {ok, TabList} | {error, Reason}
%% FileBasename = string()
%% TabList = [Tab | ..] 恢复的表
%% Tab = atom()
%% Args = [{key, Val}] 参数，查看mnesia:restore/2
%% Reason = term()
%% 恢复数据库
restore(FileBasename) ->
    restore(FileBasename, []).
restore(FileBasename, Args) ->
    Mod = env_mgr:get(?ENV_MNESIA_HANDLER),

    OptSwitch = fun
                    ({ram_copies, []}) -> {ram_copies, []};
                    ({disc_copies, []}) -> {disc_copies, []};
                    ({disc_only_copies, []}) -> {disc_only_copies, []};
                    ({ram_copies, _}) -> {ram_copies, [node()]};
                    ({disc_copies, _}) -> {disc_copies, [node()]};
                    ({disc_only_copies, _}) -> {disc_only_copies, [node()]};
                    (Other) -> Other
                end,

    Convert = fun
                  ({schema, db_nodes, _Nodes}, Acc) ->
                      {[{schema, db_nodes, [node()]}], Acc};
                  ({schema, version, Version}, Acc) ->
                      {[{schema, version, Version}], Acc};
                  ({schema, cookie, Cookie}, Acc) ->
                      {[{schema, cookie, Cookie}], Acc};
                  ({schema, Tab, CreateList}, Acc) ->
                      {[{schema, Tab, lists:map(OptSwitch, CreateList)}], Acc};
                  (Other, Acc) ->
                      {[Other], Acc}
              end,

    DirVar = env_mgr:get(?ENV_DIR_VAR),
    Path = lists:concat([DirVar, FileBasename, ".mb"]),
    PathTo = lists:concat([DirVar, FileBasename, ".restore.mb"]),
    mnesia:traverse_backup(Path, PathTo, Convert, switched),

    SkipTables = [Tab || #mnesia_table{table = Tab, copies = ram_copies} <- Mod:get_tables()],
    case mnesia:restore(PathTo, [{skip_tables, SkipTables}, {default_op, recreate_tables} | Args]) of
        {atomic, TabList} ->
            ok = file:delete(PathTo),
            {ok, TabList};
        {aborted, Reason} -> {error, Reason}
    end.
